package com.ailanyin.common.security.filter;

import cn.hutool.core.io.IoUtil;
import cn.hutool.extra.servlet.ServletUtil;
import cn.hutool.json.JSONUtil;
import com.ailanyin.common.base.result.AjaxResult;
import com.ailanyin.common.redis.service.RedisService;
import com.ailanyin.common.security.constants.SecurityConstant;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * 接口限流过滤器
 *
 * @author ailanyin
 * @version 1.0
 * @since 2022/6/9 0009 上午 9:12
 */
@Component
public class RateLimitFilter extends OncePerRequestFilter {

    @Autowired
    private RedisService redisService;

    private void returnError(ServletResponse servletResponse) throws IOException {
        //解决前端跨域
        HttpServletResponse response = (HttpServletResponse) servletResponse;
        response.setHeader("Access-Control-Allow-Origin" , "*");
        response.setHeader("Access-Control-Allow-Methods" , "POST, GET, OPTIONS, DELETE");
        response.setHeader("Access-Control-Max-Age" , "3600");
        response.setHeader("Access-Control-Allow-Headers" , "x-requested-with");
        servletResponse.setCharacterEncoding("UTF-8");
        servletResponse.setContentType("application/json; charset=utf-8");
        PrintWriter out = servletResponse.getWriter();
        out.append(JSONUtil.toJsonStr(AjaxResult.rateLimit()));
        IoUtil.close(out);
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain chain) throws ServletException, IOException {
        if (rateLimitHandle(request, response)) {
            returnError(response);
            return;
        }
        chain.doFilter(request, response);
    }

    private boolean checkUriNeedHandle(HttpServletRequest request) {
        return SecurityConstant.LIMIT_URI.contains(request.getRequestURI());
    }

    private boolean rateLimitHandle(HttpServletRequest request, ServletResponse response) throws IOException {
        if (!checkUriNeedHandle(request)) {
            return false;
        }
        String ip = ServletUtil.getClientIP(request);
        Integer limit = 0;
        if (redisService.hasKey(SecurityConstant.RATE_LIMIT_KEY + ip)) {
            limit = (Integer) redisService.get(SecurityConstant.RATE_LIMIT_KEY + ip);
            if (limit >= SecurityConstant.MAX_RATE_LIMIT_TOTAL) {
                return true;
            }

        }
        limit++;
        redisService.set(SecurityConstant.RATE_LIMIT_KEY + ip, limit, 1);
        return false;
    }
}
